#!/usr/bin/env python
# encoding: utf-8

import argparse
import torch
import numpy as np

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_data_list_path', type=str)
    parser.add_argument('--output_sub_list_path', type=str)
    parser.add_argument('--size', type=int, default=10000)
    args = parser.parse_args()

    audioSequences, Speakers = torch.load(args.input_data_list_path)
    np.random.shuffle(audioSequences)
    audioSequences = audioSequences[:args.size]

    print(f'Saved data list file at {args.output_sub_list_path}')
    torch.save((audioSequences, Speakers), args.output_sub_list_path)

